from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from .database import get_db
from .models import *
from sqlalchemy import func
import akshare as ak
from datetime import datetime, timedelta
import math
import pandas as pd
import json
from fastapi.responses import JSONResponse

router = APIRouter()


stock_cache = {
    "stock_zh_a_spot_em": {
        "interval": 7,
        "last_updated": datetime.strptime("2024-11-06", "%Y-%m-%d"),
    }
}


@router.get("/workday")
def read_workday(offset: int = 7, db: Session = Depends(get_db)):
    # Get today's date
    today = datetime.today().date()

    # Calculate the date range
    start_date = today - timedelta(days=offset)
    end_date = today + timedelta(days=offset)

    # Query workdays in the date range
    workdaylist = (
        db.query(Workday)
        .filter(Workday.trade_date >= start_date, Workday.trade_date <= end_date)
        .all()
    )

    # If no workdays are found, fetch from the external source
    if not workdaylist:
        # Fetch trade dates from Akshare
        tool_trade_date_hist_sina_df = ak.tool_trade_date_hist_sina()

        # Save the trade dates to the database
        for trade_date in tool_trade_date_hist_sina_df["trade_date"]:
            # Assuming that the trade_date is a string in the format YYYYMMDD
            workday = Workday(trade_date=trade_date)
            db.add(workday)

        # Commit the transaction to save the data
        db.commit()

        # After saving to the DB, query again for the filtered workday list
        workdaylist = (
            db.query(Workday)
            .filter(Workday.trade_date >= start_date, Workday.trade_date <= end_date)
            .all()
        )

    # Return the workday list
    return workdaylist


# 获取单个股票信息
@router.get("/stock_info")
def read_stock_info(symbol: str = "000001", db: Session = Depends(get_db)):
    stock_info = db.query(StockInfo).filter(StockInfo.symbol == symbol).first()
    if stock_info is None:
        raise HTTPException(status_code=404, detail="Stock not found")
    return stock_info


@router.get("/stock_info")
def read_all_stock_info(db: Session = Depends(get_db)):
    current_time = datetime.now()
    interval = timedelta(days=stock_cache["stock_zh_a_spot_em"]["interval"])

    stocklist = db.query(StockInfo).all()
    if (
        stock_cache["stock_zh_a_spot_em"]["last_updated"] is None
        or current_time - stock_cache["stock_zh_a_spot_em"]["last_updated"] > interval
    ):
        stock_df = ak.stock_zh_a_spot_em()
        stock_df = stock_df[
            ["代码", "名称", "总市值", "流通市值", "市盈率-动态", "市净率"]
        ]
        stock_df.columns = [
            "symbol",
            "name",
            "market_cap",
            "circulating_cap",
            "pe_ratio_dynamic",
            "pb_ratio",
        ]
        for index, row in stock_df.iterrows():
            existing_stock = next(
                (stock1 for stock1 in stocklist if stock1.symbol == row["symbol"]),
                None,
            )
            try:
                if existing_stock is None:
                    stock_info_df = ak.stock_individual_info_em(symbol=row["symbol"])
                    # 从stock_info_df获取并填充new_stock 行业-sector 上市时间-listing_date 总股本-total_shares
                    name = stock_info_df.loc[
                        stock_info_df["item"] == "股票简称", "value"
                    ].values[0]
                    industry = stock_info_df.loc[
                        stock_info_df["item"] == "行业", "value"
                    ].values[0]
                    listing_date_str = stock_info_df.loc[
                        stock_info_df["item"] == "上市时间", "value"
                    ].values[0]
                    total_shares = stock_info_df.loc[
                        stock_info_df["item"] == "总股本", "value"
                    ].values[0]

                    listing_date = (
                        datetime.strptime(str(listing_date_str), "%Y%m%d").date()
                        if not str(listing_date_str) == "-"
                        else None
                    )
                    new_stock = StockInfo(
                        symbol=row["symbol"],
                        name=name,
                        market_cap=(
                            row["market_cap"]
                            if not math.isnan(row["market_cap"])
                            else None
                        ),
                        circulating_cap=(
                            row["circulating_cap"]
                            if not math.isnan(row["circulating_cap"])
                            else None
                        ),
                        pe_ratio_dynamic=(
                            row["pe_ratio_dynamic"]
                            if not math.isnan(row["pe_ratio_dynamic"])
                            else None
                        ),
                        pb_ratio=(
                            row["pb_ratio"] if not math.isnan(row["pb_ratio"]) else None
                        ),
                        sector=(industry if not industry == "-" else None),
                        listing_date=listing_date,
                        total_shares=(
                            total_shares
                            if not str(total_shares) == "-"
                            and not math.isnan(total_shares)
                            else None
                        ),
                    )
                    db.add(new_stock)
                    db.commit()
                    print(f"Added {index} new stock: {new_stock}")

            except Exception as e:
                print(f"An error occurred: " + row["symbol"] + "," + row["name"])

        stock_cache["stock_zh_a_spot_em"]["last_updated"] = current_time

    return stocklist


@router.get("/sector_funds2")
def get_sector_funds2(db: Session = Depends(get_db)):
    stock_fund_df = ak.stock_sector_fund_flow_rank(sector_type="概念资金流")
    json_data = stock_fund_df.to_json(orient="records", force_ascii=False)
    return JSONResponse(content=json.loads(json_data), media_type="application/json")


# Assuming `Workday` has a `date` field with trading dates
def get_previous_workdays(db: Session, offset: int):
    # Query the last three workdays
    workdays = (
        db.query(Workday.trade_date)
        .filter(Workday.trade_date <= datetime.now().date())
        .order_by(Workday.trade_date.desc())
        .limit(offset)
        .all()
    )
    return [workday[0] for workday in workdays]


@router.get("/sector_funds")
def get_sector_funds(db: Session = Depends(get_db)):
    today = datetime.now()

    # Get the previous workdays from the database
    previous_workdays = get_previous_workdays(db, 4)
    if len(previous_workdays) < 3:
        raise ValueError("Insufficient workdays in database")

    _, yesterday, two_days_ago, three_days_ago = previous_workdays
    funds_hist_db = (
        db.query(SectorDaily).filter(SectorDaily.date >= three_days_ago).all()
    )
    stock_fund_df = ak.stock_sector_fund_flow_rank()
    # Add new columns with default None values
    stock_fund_df["昨天主力净流入-净额"] = None
    stock_fund_df["2天前主力净流入-净额"] = None
    stock_fund_df["3天前主力净流入-净额"] = None
    for index, stock_fund in stock_fund_df.iterrows():
        sector_funds_db = [
            record
            for record in funds_hist_db
            if record.name == stock_fund["名称"]
            and record.date in {yesterday, two_days_ago, three_days_ago}
        ]
        if len(sector_funds_db) < 3:
            stock_sector_fund_flow_hist_df = ak.stock_sector_fund_flow_hist(
                symbol=stock_fund["名称"]
            )
            stock_sector_fund_flow_hist_df = stock_sector_fund_flow_hist_df.rename(
                columns={
                    "日期": "date",
                    "名称": "name",
                    "主力净流入-净额": "main_net_inflow",
                    "主力净流入-净占比": "main_net_ratio",
                    "超大单净流入-净额": "huge_net_inflow",
                    "超大单净流入-净占比": "huge_net_ratio",
                    "大单净流入-净额": "large_net_inflow",
                    "大单净流入-净占比": "large_net_ratio",
                    "中单净流入-净额": "medium_net_inflow",
                    "中单净流入-净占比": "medium_net_ratio",
                    "小单净流入-净额": "small_net_inflow",
                    "小单净流入-净占比": "small_net_ratio",
                }
            )

            # Replace NaN with None to be compatible with SQL NULL values
            stock_sector_fund_flow_hist_df = stock_sector_fund_flow_hist_df.where(
                pd.notna(stock_sector_fund_flow_hist_df), None
            )

            # Set the name column value to `stock_fund["名称"]` if it's missing in the DataFrame
            stock_sector_fund_flow_hist_df["name"] = stock_fund["名称"]

            for _, row in stock_sector_fund_flow_hist_df.iterrows():
                if row["date"] in {
                    yesterday,
                    two_days_ago,
                    three_days_ago,
                } and not any(  # 只更新三天内的
                    stock for stock in sector_funds_db if stock.date == row["date"]
                ):
                    tosave = SectorDaily(**row.to_dict())
                    db.add(tosave)

                if row["date"] == yesterday:
                    stock_fund["昨天主力净流入-净额"] = row["main_net_inflow"]
                elif row["date"] == two_days_ago:
                    stock_fund["2天前主力净流入-净额"] = row["main_net_inflow"]
                elif row["date"] == three_days_ago:
                    stock_fund["3天前主力净流入-净额"] = row["main_net_inflow"]
        else:
            date_to_inflow = {
                record.date: record.main_net_inflow for record in sector_funds_db
            }
            stock_fund_df.loc[index, "昨天主力净流入-净额"] = date_to_inflow.get(
                yesterday
            )
            stock_fund_df.loc[index, "2天前主力净流入-净额"] = date_to_inflow.get(
                two_days_ago
            )
            stock_fund_df.loc[index, "3天前主力净流入-净额"] = date_to_inflow.get(
                three_days_ago
            )
    db.commit()
    json_data = stock_fund_df.to_json(orient="records", force_ascii=False)
    return JSONResponse(content=json.loads(json_data), media_type="application/json")


@router.get("/limitup")
def get_limitup(db: Session = Depends(get_db)):
    offset = 7
    today = datetime.now().date()
    now = datetime.now()
    days = get_previous_workdays(db, offset)
    # Query the stock limit up data within the last 5 days
    stockup_list = (
        db.query(StockDailyLimitUp)
        .filter(
            StockDailyLimitUp.date >= days[-1],
        )
        .all()
    )
    columnMap = {
        "代码": "symbol",  # 股票代码
        "名称": "name",  # 股票名称
        "涨跌幅": "percentage_change",  # 涨跌幅
        "最新价": "latest_price",  # 最新价
        "成交额": "amount",  # 成交额
        "流通市值": "circulating_market_value",  # 流通市值
        "总市值": "total_market_value",  # 总市值
        "换手率": "amount_rate",  # 换手率
        "封板资金": "board_funds",  # 封板资金
        "首次封板时间": "first_board_time",  # 首次封板时间
        "最后封板时间": "last_board_time",  # 最后封板时间
        "炸板次数": "board_break_count",  # 炸板次数
        "涨停统计": "limit_up_statistics",  # 涨停统计
        "连板数": "consecutive_boards",  # 连板数
        "所属行业": "sector",  # 所属行业
    }
    # Check if the current time is past 16:30 and if today's data is missing
    latest_entry_date = max([record.date for record in stockup_list], default=None)
    unique_dates = set([stock.date for stock in stockup_list])
    if now.hour > 16 and (latest_entry_date != today or len(unique_dates) < offset):
        for day in days:
            if day in unique_dates:
                continue
            # Fetch new data from Akshare and append it to the database
            stock_zt_pool_em_df = ak.stock_zt_pool_em(date=day.strftime("%Y%m%d"))
            if stock_zt_pool_em_df.empty:
                continue
            stock_zt_pool_em_df = stock_zt_pool_em_df.drop(
                columns=stock_zt_pool_em_df.columns[0]
            )

            stock_zt_pool_em_df.rename(
                columns=columnMap,
                inplace=True,
            )
            stock_zt_pool_em_df["date"] = day

            stock_zt_pool_em_df.to_sql(
                name="stock_daily_limit_up",
                con=db.get_bind(),
                if_exists="append",
                index=False,
            )
    elif now.hour < 16 and today not in unique_dates:
        stock_zt_pool_em_df_today = ak.stock_zt_pool_em(date=today.strftime("%Y%m%d"))
        stock_zt_pool_em_df_today = stock_zt_pool_em_df_today.drop(
            columns=stock_zt_pool_em_df_today.columns[0]
        )

        stock_zt_pool_em_df_today.rename(
            columns=columnMap,
            inplace=True,
        )
        stock_zt_pool_em_df_today["date"] = today
        # Re-query the data after update
        stockup_list = (
            db.query(StockDailyLimitUp)
            .filter(
                StockDailyLimitUp.date >= days[-1],
            )
            .all()
        )
        if not stock_zt_pool_em_df_today.empty:
            for _, row in stock_zt_pool_em_df_today.iterrows():
                rt_stockup = StockDailyLimitUp(**row.to_dict())
                stockup_list.append(rt_stockup)

    # 东财 龙虎榜 stock_lhb_detail_em

    return stockup_list


@router.get("/longhubang")
def get_longhubang(
    db: Session = Depends(get_db),
):
    offset = 365
    today = datetime.now().date()

    now = datetime.now()
    days = get_previous_workdays(db, offset)

    stockup_list = (
        db.query(StockDailyLongHuBang)
        .filter(
            StockDailyLongHuBang.date >= days[-1], StockDailyLongHuBang.date <= today
        )
        .all()
    )

    # Check if the current time is past 16:30 and if today's data is missing
    unique_dates = set([stock.date for stock in stockup_list])
    if (
        (now.hour > 16 and today not in unique_dates)
        or (today == days[0] and now.hour < 16 and days[1] not in unique_dates)
        or days[-1] not in unique_dates
    ):
        # Fetch new data from Akshare and append it to the database
        stock_lhb_detail_em_df = ak.stock_lhb_detail_em(
            start_date=days[-1].strftime("%Y%m%d"), end_date=days[0].strftime("%Y%m%d")
        )
        if not stock_lhb_detail_em_df.empty:
            stock_lhb_detail_em_df = stock_lhb_detail_em_df.drop(
                columns=stock_lhb_detail_em_df.columns[0]
            )
            stock_lhb_detail_em_df = stock_lhb_detail_em_df.replace(
                {pd.NA: None, pd.NaT: None, float("nan"): None}
            )
            tosavelist = []
            for index, row in stock_lhb_detail_em_df.iterrows():
                existing_record = next(
                    (
                        stock
                        for stock in stockup_list
                        if stock.symbol == row["代码"] and stock.date == row["上榜日"]
                    ),
                    None,
                )
                if not existing_record and not any(
                    stock["symbol"] == row["代码"] and stock["date"] == row["上榜日"]
                    for stock in tosavelist
                ):
                    new_stock = StockDailyLongHuBang(
                        symbol=row["代码"],
                        name=row["名称"],
                        date=row["上榜日"],
                        analysis=row["解读"],
                        closing_price=row["收盘价"],
                        percentage_change=row["涨跌幅"],
                        net_buy_amount=row["龙虎榜净买额"],
                        buy_amount=row["龙虎榜买入额"],
                        sell_amount=row["龙虎榜卖出额"],
                        longhubang_amount=(row["龙虎榜成交额"]),
                        market_total_amount=(row["市场总成交额"]),
                        net_buy_ratio=(row["净买额占总成交比"]),
                        amount_ratio=(row["成交额占总成交比"]),
                        amount_rate=(row["换手率"]),
                        circulating_market_value=(row["流通市值"]),
                        reason_for_listing=(row["上榜原因"]),
                        after_1_day=row.get("上榜后1日", None),
                        after_2_days=row.get("上榜后2日", None),
                        after_5_days=row.get("上榜后5日", None),
                        after_10_days=row.get("上榜后10日", None),
                    )

                    tosavelist.append(
                        {"date": new_stock.date, "symbol": new_stock.symbol}
                    )
                    db.add(new_stock)
                if (index + 1) % 20 == 0:
                    db.commit()

            stockup_list = (
                db.query(StockDailyLongHuBang)
                .filter(
                    StockDailyLongHuBang.date >= days[-1],
                    StockDailyLongHuBang.date <= today,
                )
                .all()
            )

    # 东财 龙虎榜 stock_lhb_detail_em

    return stockup_list
